from keras.layers import Input, Dense, Flatten, Reshape, Conv2D, UpSampling2D, MaxPooling2D
from keras.models import Model, Sequential
from keras.datasets import mnist
from keras.callbacks import Callback

import numpy as np
import wandb
from wandb.keras import WandbCallback

run = wandb.init()
config = run.config

config.epochs = 2

(x_train, _), (x_test, _) = mnist.load_data()

x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.

model = Sequential()
model.add(Reshape((28,28,1), input_shape=(28,28)))
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(2,2))
model.add(Conv2D(32, (3, 3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(1, (3, 3), activation='sigmoid', padding='same'))
model.add(Reshape((28,28)))

model.compile(optimizer='adam', loss='mse')

model.summary()

# For visualization
class Images(Callback):
      def __init__(self, validation_data):
            self.validation_data = validation_data

      def on_epoch_end(self, epoch, logs):
            indices = np.random.randint(self.validation_data[0].shape[0], size=8)
            test_data = self.validation_data[0][indices]
            pred_data = self.model.predict(test_data)
            wandb.log({
                  "examples": [
                        wandb.Image(np.hstack([data, pred_data[i]]), caption=str(i))
                        for i, data in enumerate(test_data)]},
                  step=epoch)

model.fit(x_train, x_train,
                epochs=config.epochs,
                validation_data=(x_test, x_test),
          callbacks=[Images((x_test, x_test)), WandbCallback()])


model.save('auto-cnn.h5')


